import HelperFunctions as HF
import numpy as np
from numpy import array
import csv
import random
from datetime import datetime
from time import strftime

'''
Global and Local Empirical Bayes Smoothers with Gamma Model
'''

def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def get_inputs():
    """Given the list of inputs for the script when run, return an object
    containing the variables for the script"""
    try:
        
        sInputFC = HF.pGP.GetParameterAsText(0)                     # Input Feature Class
        sField_Ob = HF.pGP.GetParameterAsText(1).upper()              # Observation Field Name
        sField_Pop = HF.pGP.GetParameterAsText(2).upper()                # Population Field Name
        sCongtiguityFC = HF.pGP.GetParameterAsText(3)                # Congtiguity File (.csv) with "ID1" and "ID2"
        sOutputFC = HF.pGP.GetParameterAsText(4)                  # Output Feature Class
        '''
        sInputFC = "C:\_WORK\MyPythonWork\CreateData\TP1000_cancer__.shp"
        sField_Ob = "SUM_CANCER"
        sField_Pop = "COUNT_"
        sCongtiguityFC = "C:\_WORK\MyPythonWork\CreateData\TP1000_cancer__Congtiguity.csv" 
        sOutputFC = "C:\_WORK\MyPythonWork\CreateData\TP1000_cancer__EBS.shp"
        '''

  
    except: raise HF.ReportError (HF.msgParseErr)
    
    #Make a generic object, populate it with our variables, then return it
    try:
        obj = HF.Gen()
        obj.add('sInputFC',sInputFC)
        obj.add('sField_Ob',sField_Ob)
        obj.add('sField_Pop', sField_Pop)
        obj.add('sCongtiguityFC', sCongtiguityFC)
        obj.add('sOutputFC', sOutputFC)

    except: raise HF.ReportError (HF.msgInputsErr)

    return obj

def build_value_lists():
    """Reads through feature table extracting geometry and field info.
       Returns this information in dictionary structures."""

    pRows = HF.pGP.SearchCursor(inputs.sInputFC)
    # pRows = pGP.SearchCursor(inputs.sOutputFC)
    pRow = pRows.Next()
    sKey = np.array([])
    lsBadRecs = []
    iBadRecCnt = 0
    iRec = 0 

    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage (HF.sBuildingLists)
    print "Strat building attribute lists..."
    
    
    while pRow <> None:
        iRec = iRec + 1
        try:
            dTemp = []
            dTemp.append(pRow.GetValue("FID"))
            dTemp.append(pRow.GetValue(inputs.sField_Ob))
            dTemp.append(pRow.GetValue(inputs.sField_Pop))
            sKey=np.append(sKey,dTemp)
        except:
            iBadRecCnt += 1
        pRow = pRows.Next()
    
    #print iBadRecCnt     
    # Report any problems encountered reading feature input data.
    if iBadRecCnt:
        sMessage = HF.msgReadErr % (iBadRecCnt, iRec)
        HF.pGP.AddWarning (sMessage)
        HF.pGP.AddWarning(`lsBadRecs`)


    sKey.shape = (-1,3)
    
    #print sKey    
    return sKey

def build_contiguity_list():
    spContiguity = np.array([],dtype=int)
    #read contiguity file
    fn = inputs.sCongtiguityFC

    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Start building contiguity lists...")
    print "Strat building contiguity lists..."
       
    for record in csv.DictReader(file(fn), dialect="excel"):
        temp = [int(record['id1']), int(record['id2'])]
        #print temp
        spContiguity=np.append(spContiguity, temp)
    spContiguity.shape=(-1,2)
    return spContiguity

def contiguity_item(ID):
    i = 0
    ROW = spContiguity.shape[0]
    #print ROW
    contiguityItem = []
    j = 0
    
    while i < ROW:
        if spContiguity[i,0]== ID:
            contiguityItem.append(spContiguity[i,1])
            j = j + 1
        i = i + 1
    if j == 1:
        print ID, "does not have contiguity items."
    #elif j == 
    return contiguityItem

def global_EBS():
    sumID, sumOb, sumPop = attri.sum(axis=0)
    print "sum of cancer is ", sumOb, ", sum of population is ",  sumPop
    aveID, aveOb, avePop = attri.mean(axis=0)
    print "mean of cancer is ", aveOb, ", mean of population is ",  avePop

    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Start calculating global empirical Bayes smoother...")
    print "Start calculating global empirical Bayes smoother..."    
    
    u = sumOb/sumPop

    iLen = attri.shape
    #print iLen
    i = 0
    
    sigma = 0
    while i < iLen[0]:
        sigma = attri[i][2]*(attri[i][1]/attri[i][2]-u)*(attri[i][1]/attri[i][2]-u) + sigma        
        i = i + 1
    sigma = sigma/sumPop-u/avePop
    #sigma = temp

    
    pi_EB = np.array([])
    
    i = 0    
    while i < iLen[0]:
        if sigma > 0:
            temp = 0
            w = sigma/(sigma + u/attri[i][2])
            temp = w*attri[i][1]/attri[i][2]+(1-w)*u
            temp = temp * 100
        else:
            temp = u*100
        pi_EB = np.append(pi_EB, temp)
        i = i + 1
    #print pi_EB
    print "============================Global EB===================================="
    print "mu = ", u, "sigma = ", sigma
    '''
    alpha = u*u/sigma
    beta = u/sigma
    pi = np.array([])
    i = 0
    while i < iLen[0]:
        temp = 0
        temp = (attri[i][1]+alpha)/(attri[i][2]+beta)
        temp = temp * 100
        pi = np.append(pi, temp)
        i = i + 1
    #print pi

    i = 0
    while i < iLen[0]:
        print pi_EB[i], pi[i]
        i = i + 1
    '''
    return pi_EB
    
def local_EBS():
    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Start calculating local empirical Bayes smoother...")
    print "Start calculating local empirical Bayes smoother..."
    
    pi_EB = np.array([])
    iLen = attri.shape
    #print iLen
    i = 0
    sigma = np.array([])
    mu = np.array([])
    cong_sumpop = np.array([])
    #cong_avepop = np.array([])
    while i < iLen[0]:
        contiguityItem = contiguity_item(i)
        #mu_temp = 0
        temp_ob = 0
        temp_pop = 0
        j = 0
        temp_ID = 0
        #print len(contiguityItem)
        while j < len(contiguityItem):
            #print "j = ",j
            temp_ID = contiguityItem[j]
            temp_ob = temp_ob + attri[temp_ID][1]
            temp_pop = temp_pop + attri[temp_ID][2]
            #temp = attri[j][2]*(attri[j][1]/attri[j][2]-u)*(attri[j][1]/attri[j][2]-u) + temp
            j = j + 1
        mu = np.append(mu, temp_ob/temp_pop)
        cong_sumpop = np.append(cong_sumpop, temp_pop)
        #cong_avepop = np.append(cong_avepop, temp_pop/len(contiguityItem))
        i = i + 1

    i = 0
    while i < iLen[0]:
        contiguityItem = contiguity_item(i)
        j = 0
        temp_sigma = 0
        while j < len(contiguityItem):
            temp_ID = contiguityItem[j]
            temp_sigma = attri[temp_ID][2]*(attri[temp_ID][1]/attri[temp_ID][2]-mu[temp_ID])*(attri[temp_ID][1]/attri[temp_ID][2]-mu[temp_ID]) + temp_sigma
            j = j + 1
        temp_sigma = temp_sigma/cong_sumpop[i] - mu[i]*len(contiguityItem)/cong_sumpop[i]
        sigma = np.append(sigma, temp_sigma)
        i = i + 1
    #print sigma[1]
    
    i = 0
    while i < iLen[0]:
        if sigma[i]<0:
            temp = mu[i]*100
        else:
            temp_weight = 0
            temp_weight = sigma[i]/(sigma[i]+mu[i]/attri[i][2])
            temp = 0
            temp = temp_weight*attri[i][1]/attri[i][2]+(1-temp_weight)*mu[i]
            temp = temp * 100
        pi_EB = np.append(pi_EB, temp)
        i = i + 1
    #print pi_EB
    return pi_EB
  

def output_results():
    """ This function writes results to the output feature class. """
    
    # Add results field.
    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Writing results fields...")
    print "Writing results fields..." 
    #addField = []
    #print "========================================================"    
    
    outField = "r_GlobalEB"
    if properties.dcFields.has_key(outField.upper()):
        HF.pGP.deletefield(inputs.sOutputFC, outField) 
    HF.pGP.AddField(inputs.sOutputFC, outField, "FLOAT")
    outField_1 = "r_LocalEB"
    if properties.dcFields.has_key(outField_1.upper()):
        HF.pGP.deletefield(inputs.sOutputFC, outField_1) 
    HF.pGP.AddField(inputs.sOutputFC, outField_1, "FLOAT")
    
    # Add results to output FC
    #HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    #HF.pGP.AddMessage (HF.sWritingResults)
    pRows = HF.pGP.UpdateCursor(inputs.sOutputFC)
    pRow = pRows.Next()
    iError = 0
    iCnt = 0
    fInterval = 10000 / 5.0
    fMore = fInterval
    iComplete = 20


    while pRow <> None:
        iKey = pRow.GetValue("FID")

        try:
            pRow.SetValue(outField, r_GlobalEB[iKey])
            pRow.SetValue(outField_1, r_LocalEB[iKey])
            pRows.UpdateRow(pRow)
            iCnt = iCnt + 1
            if iCnt > fInterval: fInterval, iComplete = HF.check_progress(fInterval, fMore, iComplete)
        
        except:
            iError = iError + 1
        
        pRow = pRows.Next()
        #print iKey
        

    #print "++++++++++++++++++++++"
    print "Error = ", iError
    
        
    HF.pGP.AddMessage (HF.s100Percent)
    HF.pGP.AddMessage(" ")
    pRows = None           
    #return sField



    
    
    
#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
    inputs = get_inputs()
    #print inputs.sInputFC
    
    if HF.pGP.exists(inputs.sOutputFC):
        HF.pGP.delete(inputs.sOutputFC)

    #Copy the input feature class to the output feature class.
    try:
        HF.pGP.QualifiedFieldNames = 0
        HF.pGP.Copyfeatures(inputs.sInputFC, inputs.sOutputFC)
    except:
        sMessage = HF.msgOverwriteErr % (inputs.sOutputFC)
        raise HF.ReportError (sMessage)
    
    properties = HF.get_featureclass_properties(inputs.sInputFC)
    attri = build_value_lists()
    spContiguity = build_contiguity_list()
    
    r_GlobalEB = global_EBS()
    r_LocalEB = local_EBS()


    output_results()  

           
